from ultralytics import YOLO

if __name__ == '__main__':
    # 直接使用预训练模型创建模型.
    # model = YOLO('yolov8n.pt')
    # model.train(**{'cfg': 'ultralytics/cfg/exp1.yaml', 'data': 'dataset/data.yaml'})

    # 使用yaml配置文件来创建模型,并导入预训练权重.
    # model = YOLO('ultralytics-main/ultralytics/cfg/models/v8/yolov8n-CoordConv.yaml')
    model = YOLO('E:/深度学习框架与应用/ultralytics-main/ultralytics-main/ultralytics/cfg/models/v8/yolov8n-CoordConv.yaml')

    # model.train(cfg="ultralytics/cfg/default.yaml",data="ultralytics/datasets/safe.yaml",
    #             epochs=50,batch=2,workers=2)
    model.train(cfg="ultralytics/cfg/default.yaml", data="E:/深度学习框架与应用/ultralytics-main/ultralytics-main/ultralytics/datasets/safe.yaml",
                epochs=1, batch=1, workers=1)



    # # 模型验证
    # model = YOLO('runs/detect/yolov8n_exp/weights/best.pt')
    # model.val(**{'data': 'dataset/data.yaml'})
    model = YOLO('runs/detect/train9/weights/best.pt')
    model.val(**{'data': 'datasets/safe.yaml'})


    #模型推理
    # model = YOLO('runs/detect/yolov8n_exp/weights/best.pt')
    # model.predict(source='dataset/images/test', **{'save': True})
    model = YOLO('runs/detect/train9/weights/best.pt')
    model.predict(source='datasets/safe/test/images', **{'save': True})